#! /usr/bin/env python
# coding=utf-8

import os

import shutil

from ete3 import PhyloTree

from Bio import SeqIO

import argparse
import sys



parser = argparse.ArgumentParser(
    description='''给定一个ogs比对文件 结果 对每个ogs 输出一个文件夹 作为工作目录
    需要给定 比对的ogs文件; 物种树; 所有物种的mRNA文件 
    注意 这里物种名字必须是缩写 以 如 Aae的形式 ogs每条序列的抬头必须用@分割物种 如Aae@rna-XM_021838739.1
    tree中的每一个页节点需要是物种， mRNA文件需要是每个物种单独一个文件夹如 Aae_mRNA.fa 里面的名字需要与ogs aa比对的一致
    用法:
    make_alignment_file_for_each_ogs.py  -a Orthogroup_Sequences_mafft_raw -m All_mRNA -t spe_tree -o OGS_for_testing
    由大天才于2021年11月10日创建于浙江农业大学''')


parser.add_argument('-t',
                help='必须给定，物种树')

parser.add_argument('-a',
                help='必须给定,比对好的ogs氨基酸文件')

parser.add_argument('-m',
                help='必须给定，所有物种的mRNA文件')

parser.add_argument('-o',
                help='必须给定,输出的路径')




args = parser.parse_args()

if not args.a or not args.t or not args.m or not args.o:
    parser.print_help()
    sys.exit()





outpath = args.o

tree_file = args.t

orthopath = args.a

mRNA_path = args.m

try:
	shutil.rmtree(outpath)
except:
	pass

try:
	os.mkdir(outpath)
except:
	pass



# 读取mRNA



mRNA_dic = {}

for i in os.listdir(mRNA_path):

    spe = i.split('_')[0]
    mRNA_dic[spe] = {}
    for j in SeqIO.parse(mRNA_path+'/'+i,'fasta'):
        mRNA_dic[spe][str(j.name)] = str(j.seq)
        
print('finish reading mRNA file')

for i in os.listdir(orthopath):

    

    ogs_name = i.split('.')[0]
    target_path = outpath+'/'+ ogs_name
    print(ogs_name)
    try:
        os.mkdir(target_path)
    except:
        pass
    #shutil.copy(tree_file, target_path+'/spe_tree')
    
    pep_raw_file = open(target_path+'/pep_raw_ali.fa','w')
    mRNA_raw_file = open(target_path+'/mRNA_raw_ali.fa','w')
    tree_spe = open(target_path+'/spe_tree','w')
    
    # 开始读取数据
    pep_lista = []
    mRNA_lista = []
    spe_lista = []
    for j in SeqIO.parse(orthopath+'/'+i,'fasta'):
        name = str(j.name)
        spe = name.split('@')[0]
        spe_lista.append(spe)
        
        pep_raw = str(j.seq)
        pep_lista.append(pep_raw)
        
        mRNA_raw = mRNA_dic[spe][name]
        
        mRNA_ali = ''
        n = 0
        for ttt in pep_raw:
            if ttt=='-':
                mRNA_ali+='---'
            else:
                mRNA_ali+=mRNA_raw[n:n+3]
                n+=3
                
        mRNA_lista.append(mRNA_ali)
        
        pep_raw_file.write('>'+spe+'\n'+pep_raw+'\n')
        mRNA_raw_file.write('>'+spe+'\n'+mRNA_ali+'\n')
    # 建立物种树
    tree = PhyloTree(tree_file)
    tree.prune(spe_lista)
    tree_spe.write(tree.write(format=9)) # 支长全部清零
    pep_raw_file.close()
    mRNA_raw_file.close()
    tree_spe.close()